#include"effective_ensemble.hpp"

effective_ensemble::effective_ensemble(const effective_ensemble& eff){
  states=eff.states;
  //prob=eff.prob;
  occ=eff.occ;
  ntot=eff.ntot;
}

//effective_ensemble::effective_ensemble(int no_of_states,int dim_of_H){
//    states.resize(dim_of_H,no_of_states);
//    prob.resize(no_of_states);
//    occ.resize(no_of_states);
//}

effective_ensemble::effective_ensemble(matrix<std::complex<double> > mat,
				       vector<int> o){
  std::pair<std::map<int,vector<std::complex<double> > >::iterator,bool> ret_s;
  std::pair<std::map<int,int>::iterator,bool> ret_o;
  //check
  if(mat.size2()!=o.size()){
    std::cerr<<"Number of states and occurence does not match"<<std::endl;
    exit(1);
  }

  //set states
  for(unsigned jj=0;jj<mat.size2();jj++){
    vector<std::complex<double> > v=column(mat,jj);
    ret_s=states.insert(std::make_pair(jj,v));
    if(ret_s.second==false){
      states.erase(jj);
      states.insert(std::make_pair(jj,column(mat,jj)));
    }
  }
  //set occurence
  for(unsigned kk=0;kk<o.size();kk++){
    ret_o=occ.insert(std::make_pair(kk,o(kk)));
    if(ret_o.second==false){
      states.erase(kk);
      states.insert(std::make_pair(kk,o(kk)));
    }
  }
  //set total number of states
  ntot=0;
  std::map<int,int>::iterator it;
  for(it=occ.begin();it!=occ.end();it++){
    ntot += it->second;
  }
}

effective_ensemble& effective_ensemble::operator=(const effective_ensemble eff){
  states=eff.states;
  occ=eff.occ;
  ntot=eff.ntot;
  return *this;
}

vector<std::complex<double> > effective_ensemble::get_state(int alpha){
    return states[alpha];
}

void effective_ensemble::set_state(int alpha,vector<std::complex<double> > vec){
  
  std::pair<std::map<int,vector<std::complex<double> > >::iterator,bool> ret_s;
  ret_s=states.insert(std::make_pair(alpha,vec));
  
  if(ret_s.second==false){
    states.erase(alpha);
    states.insert(std::make_pair(alpha,vec));
  }
  
}
double effective_ensemble::P(int alpha){
  return double(occ[alpha])/double(ntot);
}

vector<double> effective_ensemble::P(){
  vector<double> vec(occ.size());
  std::map<int,int>::iterator it;
  for(it=occ.begin();it!=occ.end();it++){
    vec(it->first)=double(it->second)/double(ntot);
  }
  return vec;
}


matrix<std::complex<double> > effective_ensemble::rho(){
  vector<std::complex<double> > w=states[0];
  vector<std::complex<double> > wc=boost::numeric::ublas::conj(w);
  matrix<std::complex<double> > r=P(0)*outer_prod(w,wc);
  std::map<int,int>::iterator it=occ.begin();
  //move iterator once
  it++;
  for( ;it!=occ.end();it++){
    w=states[it->first];
    wc=boost::numeric::ublas::conj(w);
    r += P(it->first)*outer_prod(w,w);
  }
  return r;
}

int effective_ensemble::get_number_of_states(){
  return states.size();
}

vector<std::complex<double> > effective_ensemble::rho_vec(){
  unsigned d;
  if(!states.empty()){
    d=states[0].size();
  }
  if(states.empty()){
    std::cerr<<" effective ensemble is not initialized yet"<<std::endl;
    exit(1);
  }
  vector<std::complex<double> > vec(d);
  matrix<std::complex<double> > r=rho();
  for(unsigned ii=0;ii<d;ii++){
    for(unsigned jj=0;jj<d;jj++){
      vec(ii*d+jj)=r(ii,jj);
    }
  }
  return vec;
}
  
int effective_ensemble::get_dim(){
  int d;
  if(!states.empty()){
    d=states[0].size();
  }
  if(states.empty()){
    std::cerr<<" effective ensemble is not initialized yet"<<std::endl;
    exit(1);
  }
  return d;
}

int effective_ensemble::occupancy(int alpha){
  std::map<int,int>::iterator it;
  it=occ.find(alpha);
  if(it==occ.end()){
    std::cerr<<"Occupancy of state with key "<< alpha<<" is not found!"<<std::endl;
    exit(1);
  }
  return it->second;
}

void effective_ensemble::set_occupancy(int alpha,int n){
  
  std::pair<std::map<int,int>::iterator,bool> ret_o;
  ret_o=occ.insert(std::make_pair(alpha,n));
  
  if(ret_o.second==false){
    states.erase(alpha);
    states.insert(std::make_pair(alpha,n));
  }
}

vector<double> effective_ensemble::variance(){
  int num_states=this->get_number_of_states();
  int dim=this->get_dim();
  int N=ntot;
  identity_matrix<double> I(dim);
  matrix<std::complex<double> > dens=this->rho();
  vector<double> s(dim*dim,0);
  
  for(int ii=0;ii<dim;ii++){
    for(int jj=0;jj<dim;jj++){
      for(int kk=0;kk<num_states;kk++){
	s(ii*dim+jj)=+std::pow(std::abs((std::pow(norm_2(prod(outer_prod(column(I,ii),column(I,jj)),
							      this->get_state(kk))),2)*this->P(kk))-
					dens(ii,jj)/double(N)),2);
      }
    }
  }
  double d=double(N)/double(N-1);
  for(int ii=0;ii<dim*dim;ii++){
    s(ii)=std::sqrt(d*s(ii));
  }
  return s;
}


